#include"StringPool.h"
#include"variable.h"
#include<vector>
#include<string>
#include<fstream>
#include<ctime>
#include"Util.h"
#include"Parameters.h"
#include"DependencyInstance.h"
#include"KBestParseForest.h"
#include"DependencyDecoder.h"
using namespace std;

DependencyDecoder::DependencyDecoder(DependencyPipe *pipe){
	this->pipe = pipe;
}

DependencyDecoder::~DependencyDecoder(){}

vector<vector<int> >* DependencyDecoder::getTypes(std::vector<vector<vector<vector<double> > > > &nt_probs, int len){
	vector<vector<int> >* static_types = new vector<vector<int> >(len);
	for(int i = 0; i < len; i++){
		(*static_types)[i] = vector<int>(len);
	}

	for(int i = 0; i < len; i++){
		for(int j = 0; j < len; j++){
			if(i == j){
				(*static_types)[i][j] = 0;
				continue;
			}
			int wh = -1;
			double best = Negative_Infinity;
			int length = (int)(pipe->types.size());
			for(int t = 0; t < length; t++){
				double score = 0.0;
				if(i < j){
					score = nt_probs[i][t][0][1] + nt_probs[j][t][0][0];
				}
				else{
					score = nt_probs[i][t][1][1] + nt_probs[j][t][1][0];
				}
				if(score > best){
					wh = t;
					best = score;
				}
			}
			(*static_types)[i][j] = wh;
		}
	}
	return static_types;
}

vector<pair<FeatureVector*, string*> >* DependencyDecoder::decodeProjective(DependencyInstance *inst, std::vector<vector<vector<double> > > &probs, std::vector<vector<vector<vector<double> > > > &nt_probs, int K){
	vector<string*>* forms = inst->forms;

	int length = (int)(forms->size());
	vector<vector<int> >* static_types = NULL;
	if(pipe->labeled){
		static_types = getTypes(nt_probs, length);
	}

	KBestParseForest pf = KBestParseForest(0, length - 1, *inst, K);

	bool added = false;
	bool added1 = false, added2 = false;
	for(int s = 0; s < length; s++){
		pf.add(s, -1, 0, 0.0);
		pf.add(s, -1, 1, 0.0);
	}
	for(int j = 1; j < length; j++){
		for(int s = 0; s < length && s + j < length; s++){
			int t = s + j;
			double prodProb_st = probs[s][t][0];
			double prodProb_ts = probs[s][t][1];

			int type1 = pipe->labeled ? (*static_types)[s][t] : 0;
			int type2 = pipe->labeled ? (*static_types)[t][s] : 0;

			double nt_prob_s_01 = 0.0;
			double nt_prob_s_10 = 0.0;
			double nt_prob_t_00 = 0.0;
			double nt_prob_t_11 = 0.0;
			
			
			if(pipe->labeled){
				nt_prob_s_01 = nt_probs[s][type1][0][1];
				nt_prob_s_10 = nt_probs[s][type2][1][0];
				nt_prob_t_00 = nt_probs[t][type1][0][0];
				nt_prob_t_11 = nt_probs[t][type2][1][1];
			}
			for(int r = s; r < t; r++){
				// first is direction, second is complete
				// _s means s is the parent
				
				ParseForestItem* b1 = pf.getItems(s, r, 0, 1);
				ParseForestItem* c1 = pf.getItems(r + 1, t, 1, 1);
				vector<BestPair>* pairs = pf.getKBestPairs(b1, c1);
				int length1 = (int)(pairs->size());
				for(int k = 0; k < length1; k++){
					if(((*pairs)[k]).first == -1 || ((*pairs)[k]).second == -1){
						break;
					}

					int comp1 = (*pairs)[k].first;
					int comp2 = (*pairs)[k].second;

					double bc = b1[comp1].prob + c1[comp2].prob;

					double prob_fin = bc + prodProb_st;

					if(pipe->labeled){
						prob_fin += nt_prob_s_01 + nt_prob_t_00;
					}

					added1 = pf.add(s, r, t, type1, 0, 0, prob_fin, &(b1[comp1]), &(c1[comp2]));

					prob_fin = bc + prodProb_ts;

					if(pipe->labeled){
						prob_fin += nt_prob_t_11 + nt_prob_s_10;
					}

					added2 = pf.add(s, r, t, type2, 1, 0, prob_fin, &(b1[comp1]), &(c1[comp2]));
					if(!added1 && !added2){
						break;
					}
				}
				delete(pairs);
			}
			for(int r = s; r <= t; r++){
				if(r != s){
					ParseForestItem* b1 = pf.getItems(s, r, 0, 0);
					ParseForestItem* c1 = pf.getItems(r, t, 0, 1);
					vector<BestPair>* pairs = pf.getKBestPairs(b1, c1);
					int length1 = (int)(pairs->size());
					for(int k = 0; k < length1; k++){
						if(((*pairs)[k]).first == -1 || ((*pairs)[k]).second == -1){
							break;
						}

						int comp1 = (*pairs)[k].first;
						int comp2 = (*pairs)[k].second;

						double bc = b1[comp1].prob + c1[comp2].prob;
						added = pf.add(s, r, t, -1, 0, 1, bc, &(b1[comp1]), &(c1[comp2]));
						if(!added){
							break;
						}
					}
					delete(pairs);
				}
				if(r != t){
					ParseForestItem* b1 = pf.getItems(s, r, 1, 1);
					ParseForestItem* c1 = pf.getItems(r, t, 1, 0);
					vector<BestPair>* pairs = pf.getKBestPairs(b1, c1);
					int length1 = (int)(pairs->size());
					for(int k = 0; k < length1; k++){
						if(((*pairs)[k]).first == -1 || ((*pairs)[k]).second == -1){
							break;
						}

						int comp1 = (*pairs)[k].first;
						int comp2 = (*pairs)[k].second;

						double bc = b1[comp1].prob + c1[comp2].prob;
						added = pf.add(s, r, t, -1, 1, 1, bc, &(b1[comp1]), &(c1[comp2]));
						if(!added){
							break;
						}
					}
					delete(pairs);
				}
			}
		}
	}
	if(static_types != NULL){
		delete(static_types);
	}

	vector<pair<FeatureVector*, string*> >* result = NULL;
	if(pipe->options->root_type == 0){
		result = pf.getBestParses(inst, pipe);
	}
	else if(pipe->options->root_type == 1){
		result = pf.getBestParses(inst, pipe, probs, nt_probs, pipe->labeled);
	}
	return result;
}

vector<pair<FeatureVector*, string*> >* DependencyDecoder::decodeNonProjective(DependencyInstance *inst, std::vector<vector<vector<FeatureVector> > > &fvs, std::vector<vector<vector<double> > > &probs, std::vector<vector<vector<vector<FeatureVector> > > > &nt_fvs, std::vector<vector<vector<vector<double> > > > &nt_probs, int K){
	int numWords = inst->length();

	int *oldI = new int[numWords * numWords];
	int *oldO = new int[numWords * numWords];
	double *scoreMatrix = new double[numWords * numWords];
	double *orig_scoreMatrix = new double[numWords * numWords];
	bool *curr_nodes = new bool[numWords];
	IntHashMap *reps = new IntHashMap[numWords];

	vector<vector<int> > *static_types = NULL;
	if(pipe->labeled){
		static_types = getTypes(nt_probs, numWords);
	}
	
	for (int i = 0; i < numWords; i++){
		curr_nodes[i] = true;
		reps[i][i] = 0;
		for (int j = 0; j < numWords; j++){
			// score of edge (i,j) i --> j
			scoreMatrix[i * numWords + j] = probs[i < j ? i : j][i < j ? j : i][i < j ? 0 : 1]
			+ (pipe->labeled ? (nt_probs[i][(*static_types)[i][j]][i < j ? 0 : 1][1] + nt_probs[j][(*static_types)[i][j]][i < j ? 0 : 1][0]) : 0.0);
				
			orig_scoreMatrix[i * numWords + j] = scoreMatrix[i * numWords + j];
				
			oldI[i * numWords + j] = i;
			oldO[i * numWords + j] = j;
			if (i == j || j == 0){
				continue; // no self loops of i --> 0
			}
		}
	}

	IntHashMap final_edges = IntHashMap();
	chuLiuEdmonds(scoreMatrix, curr_nodes, oldI, oldO, false, final_edges, reps, numWords);

	int *par = new int[numWords];
	IntHashMap::iterator iter;
	for(iter = final_edges.begin(); iter != final_edges.end(); iter++){
		int ch = iter->first;
		int pr = iter->second;
		par[ch] = pr;
	}

	int *n_par = getKChanges(par, orig_scoreMatrix, (K < numWords ? K : numWords), numWords);

	delete[](oldI);
	delete[](oldO);
	delete[](scoreMatrix);
	delete[](orig_scoreMatrix);
	delete[](curr_nodes);
	delete[](reps);

	int new_k = 1;
	for (int i = 0; i < numWords; i++){
		if (n_par[i] > -1){
			new_k++;
		}
	}

	// Create Feature Vectors;
	int *fin_par = new int[new_k * numWords];
	for(int i = 0; i < numWords; i++){
		fin_par[i] = par[i];
	}
	int c = 1;
	for(int i = 0; i < numWords; i++){
		if(n_par[i] > -1){
			for(int j = 0; j < numWords; j++){
				fin_par[c * numWords + j] = par[j];
			}
			fin_par[c * numWords + i] = n_par[i];
			c++;
		}
	}
	delete[](par);
	delete[](n_par);

	vector<vector<FeatureVector*> > fin_fv = vector<vector<FeatureVector*> >(new_k, vector<FeatureVector*>(numWords));
	for(int k = 0; k < new_k; k++){
		for(int i = 0; i < numWords; i++){
			int ch = i;
			int pr = fin_par[k * numWords + i];
			if(pr != -1){
				fin_fv[k][ch] = &(fvs[ch < pr ? ch : pr][ch < pr ? pr : ch][ch < pr ? 1 : 0]);
				if(pipe->labeled){
					fin_fv[k][ch]->cat(&(nt_fvs[ch][(*static_types)[pr][ch]][ch < pr ? 1 : 0][0]),
						&(nt_fvs[pr][(*static_types)[pr][ch]][ch < pr ? 1 : 0][1]));
				}
			}
			else{
				fin_fv[k][ch] = NULL;
			}
		}
	}
	vector<pair<FeatureVector*, string*> >* d = new vector<pair<FeatureVector*, string*> >(new_k);
	for(int k = 0; k < new_k; k++){
		FeatureVector* fv = new FeatureVector();
		string* str = pool->get();
		for(int i = 0; i < numWords; i++){
			if(fin_fv[k][i] != NULL){
				fv->extend(fin_fv[k][i]);
			}
			string *ss = Util::intToString(fin_par[k * numWords + i]);
			str->append(*ss);
			pool->push_back(ss);
			str->append("|");
			ss = Util::intToString(i);
			str->append(*ss + ":");
			pool->push_back(ss);
			if(pipe->labeled){
				ss = Util::intToString((*static_types)[fin_par[k * numWords + i]][i]);
			}
			else{
				string tmp = string("0");
				ss = pool->get(tmp);
			}
			str->append(*ss + " ");
			pool->push_back(ss);
		}
		(*d)[k].first = fv;
		(*d)[k].second = Util::trim(str);
		pool->push_back(str);
	}
	delete[](fin_par);
	if(pipe->labeled){
		delete(static_types);
	}
	return d;
}

int* DependencyDecoder::getKChanges(int *par, double *scoreMatrix, int K, int numWords){
	int *result = new int[numWords];
	int *n_par = new int[numWords];
	double *n_score = new double[numWords];
	for(int i = 0; i < numWords; i++){
		result[i] = -1;
		n_par[i] = -1;
		n_score[i] = Negative_Infinity;
	}
	
	bool *isChild = calcChilds(par, numWords);

	for(int i = 0; i < numWords; i++){
		double max = Negative_Infinity;
		int wh = -1;
		for(int j = 0; j < numWords; j++){
			if(i == j || par[i] == j || isChild[i * numWords + j]){
				continue;
			}
			if(scoreMatrix[j * numWords + i] > max){
				max = scoreMatrix[j * numWords + i];
				wh = j;
			}
		}
		n_par[i] = wh;
		n_score[i] = max;
	}
	delete[](isChild);

	for(int k = 0; k < K; k++){
		double max = Negative_Infinity;
		int wh = -1, whI = -1;
		for(int i = 0; i < numWords; i++){
			if(n_par[i] == -1){
				continue;
			}
			double score = scoreMatrix[n_par[i] * numWords + i];
			if(score > max){
				max = score;
				whI = i;
				wh = n_par[i];
			}
		}

		if(max <= Negative_Infinity){
			break;
		}
		result[whI] = wh;
		n_par[whI] = -1;
	}
	delete[](n_par);
	delete[](n_score);

	return result;
}

bool* DependencyDecoder::calcChilds(int *par, int numWords){
	bool *isChild = new bool[numWords * numWords];
	memset(isChild, 0 ,sizeof(bool) * numWords * numWords);
	for(int i = 0; i < numWords; i++){
		int l = par[i];
		while(l != -1){
			isChild[l * numWords + i] = true;
			l = par[l];
		}
	}
	return isChild;
}

void DependencyDecoder::chuLiuEdmonds(double *scoreMatrix, bool *curr_nodes, int *oldI, int *oldO, bool print, IntHashMap &final_edges, IntHashMap *reps, int numWords){
	// need to construct for each node list of nodes they represent (here
	// only!)
	int *par = new int[numWords];
	// create best graph
	par[0] = -1;
	for(int i = 0; i < numWords; i++){
		// only interested in current nodes
		if(!curr_nodes[i]){
			continue;
		}
		double maxScore = scoreMatrix[i];
		par[i] = 0;
		for(int j = 0; j < numWords; j++){
			if(i == j || !curr_nodes[j]){
				continue;
			}
			double newScore = scoreMatrix[j * numWords + i];
			if(newScore > maxScore){
				maxScore = newScore;
				par[i] = j;
			}
		}
	}

	if(print){
		cout<<" After init"<<endl;
		for(int i = 0; i < numWords; i++){
			if(curr_nodes[i]){
				cout<<par[i]<<"|"<<i<<" ";
			}
		}
		cout<<endl;
	}

	// Find a cycle
	IntHashMap cycle = IntHashMap();
	bool *added = new bool[numWords];
	bool flag = false;
	memset(added, 0, sizeof(bool) * numWords);
	for(int i = 0; i <numWords && !flag; i++){
		// if I have already considered this or
		// This is not a valid node (i.e. has been contracted)
		if(added[i] || !curr_nodes[i]){
			continue;
		}
		added[i] = true;
		IntHashMap cc = IntHashMap();
		cc[i] = 0;
		int l = i;
		while(true){
			if(par[l] == -1){
				added[l] = true;
				break;
			}
			if(cycle.find(par[l]) != cycle.end()){
				cc = IntHashMap();
				int lorg = par[l];
				cc[lorg] = par[lorg];
				added[lorg] = true;
				int ll = par[lorg];
				while(ll != lorg){
					cc[ll] = par[ll];
					added[ll] = true;
					ll = par[ll];
				}
				flag = true;
				cycle = cc;
				break;
			}
			cc[l] = 0;
			l = par[l];
			if(added[l] && cc.find(par[l]) == cc.end()){
				break;
			}
			added[l] = true;
		}
	}
	delete[](added);

	// get all edges and return them
	if(!flag){
		for(int i = 0; i <numWords; i++){
			if(!curr_nodes[i]){
				continue;
			}
			if (par[i] != -1) {
				int pr = oldI[par[i] * numWords + i];
				int ch = oldO[par[i] * numWords + i];
				final_edges[ch] = pr;
			}
			else{
				final_edges[0] = -1;
			}
		}
		return;
	}

	int cycle_size = (int)(cycle.size());
	int* cyc_nodes = new int[cycle_size];
	IntHashMap::iterator iter;
	int numIter = 0;
	for(iter = cycle.begin(); iter != cycle.end(); iter++){
		cyc_nodes[numIter++] = iter->first;
	}
	int rep = cyc_nodes[0];

	if (print) {
		cout<<"Found Cycle"<<endl;
		for (int i = 0; i < cycle_size; i++){
			cout<<cyc_nodes[i]<<" ";
		}
		cout<<endl;
	}

	double cyc_weight = 0.0;
	for(int i = 0; i < cycle_size; i++){
		cyc_weight += scoreMatrix[par[cyc_nodes[i]] * numWords + cyc_nodes[i]];
	}

	for (int i = 0; i < numWords; i++){

		if (!curr_nodes[i] || cycle.find(i) != cycle.end()){
			continue;
		}

		double max1 = Negative_Infinity;
		int wh1 = -1;
		double max2 = Negative_Infinity;
		int wh2 = -1;

		for (int j = 0; j < cycle_size; j++){
			int j1 = cyc_nodes[j];

			if (scoreMatrix[j1 * numWords + i] > max1){
				max1 = scoreMatrix[j1 * numWords + i];
				wh1 = j1;
			}

			// cycle weight + new edge - removal of old
			double scr = cyc_weight + scoreMatrix[i * numWords + j1] - scoreMatrix[par[j1] * numWords + j1];
			if (scr > max2){
				max2 = scr;
				wh2 = j1;
			}
		}

		scoreMatrix[rep * numWords + i] = max1;
		oldI[rep * numWords + i] = oldI[wh1 * numWords + i];
		oldO[rep * numWords + i] = oldO[wh1 * numWords + i];
		scoreMatrix[i * numWords + rep] = max2;
		oldO[i * numWords + rep] = oldO[i * numWords + wh2];
		oldI[i * numWords + rep] = oldI[i * numWords + wh2];
	}

	IntHashMap* rep_cons = new IntHashMap[cycle_size];
	for (int i = 0; i < cycle_size; i++){
		if (print){
			cout<<cyc_nodes[i]<<": ";
		}
		IntHashMap::iterator iter;
		for (iter = reps[cyc_nodes[i]].begin(); iter != reps[cyc_nodes[i]].end(); iter++){
			rep_cons[i].insert(pair<int, int>(iter->first, 0));
			if (print){
				cout<<iter->first<<" ";
			}
		}
		if(print){
			cout<<endl;
		}
	}

	// don't consider not representative nodes
	// these nodes have been folded
	for (int i = 1; i < cycle_size; i++){
		curr_nodes[cyc_nodes[i]] = false;
		IntHashMap::iterator iter;
		for (iter = reps[cyc_nodes[i]].begin(); iter != reps[cyc_nodes[i]].end(); iter++){
			reps[rep][iter->first] = 0;
		}
	}

	chuLiuEdmonds(scoreMatrix, curr_nodes, oldI, oldO, print, final_edges, reps, numWords);

	// check each node in cycle, if one of its representatives
	// is a key in the final_edges, it is the one.
	int wh = -1;
	bool found = false;

	for(int i = 0; i < cycle_size && !found; i++){
		IntHashMap::iterator iter;
		for(iter = rep_cons[i].begin(); iter != rep_cons[i].end(); iter++){
			if(final_edges.find(iter->first) != final_edges.end()){
				wh = cyc_nodes[i];
				found = true;
				break;
			}
		}
	}
	delete[](rep_cons);

	int l = par[wh];
	while (l != wh) {
		int ch = oldO[par[l] * numWords + l];
		int pr = oldI[par[l] * numWords + l];
		final_edges[ch] = pr;
		l = par[l];
	}

	if (print){
		IntHashMap::iterator iter;
		for (iter = final_edges.begin(); iter != final_edges.end(); iter++){
			cout<<iter->second<<"|"<<iter->first<<" ";
		}
		cout<<endl;
	}
	delete[](cyc_nodes);
	delete[](par);

	return;
}
